# 将数据集划分为训练集和验证集

import os
import open3d as o3d
import numpy as np
import random
import json

import sys
import shutil
from tqdm import tqdm
# sourcery skip: avoid-builtin-shadow
if __name__ == "__main__":
    folder = r"E:\pose\datasets\obj_ac\003"
    _folder = r"E:\pose\datasets\obj_ac\obj_000003"

    train_idx = 0
    test_idx = 0
    for filename in tqdm(os.listdir(os.path.join(folder,'label_images'))):
        # rgb
        id = filename.split('.')[0].split('-')[-1]
        # print(id)

        if random.random() < 0.8:

            to_folder = os.path.join(_folder, 'train')
            if not os.path.exists(to_folder):
                os.makedirs(to_folder)
                os.makedirs(os.path.join(to_folder, 'rgb'))
                os.makedirs(os.path.join(to_folder, 'depth'))
                os.makedirs(os.path.join(to_folder, 'pose'))
                os.makedirs(os.path.join(to_folder, 'labels'))
                os.makedirs(os.path.join(to_folder, 'label_images'))

            # depth
            # shutil.copyfile((os.path.join(folder,"depth",f'{id}.jpg')),(os.path.join(to_folder,"depth",f"{train_idx}.jpg")))
            # pose

            shutil.copyfile((os.path.join(folder,"pose",f"{id}.txt")),(os.path.join(to_folder,"pose",f"{train_idx}.txt")))
            # rgb
            shutil.copyfile(
                os.path.join(folder, 'rgb', f'{id}.jpg'),
                os.path.join(to_folder, 'rgb', f'{train_idx}.jpg'),
            )

            # label
            shutil.copyfile(
                os.path.join(folder, 'labels', f'{id}.txt'),
                os.path.join(to_folder, 'labels', f'{train_idx}.txt'),
            )

            # label_images
            shutil.copyfile(
                os.path.join(folder, 'label_images', filename),
                os.path.join(to_folder, 'label_images', f'{train_idx}.jpg'),
            )

            train_idx += 1
        else:
            to_folder = os.path.join(_folder, 'val')
            if not os.path.exists(to_folder):
                os.makedirs(to_folder)
                os.makedirs(os.path.join(to_folder, 'rgb'))
                os.makedirs(os.path.join(to_folder, 'depth'))
                os.makedirs(os.path.join(to_folder, 'pose'))
                os.makedirs(os.path.join(to_folder, 'labels'))
                os.makedirs(os.path.join(to_folder, 'label_images'))
                

            # depth
            # shutil.copyfile((os.path.join(folder,"depth",f'{id}.jpg')),(os.path.join(to_folder,"depth",f"{test_idx}.jpg")))
            # pose
            shutil.copyfile((os.path.join(folder,"pose",f"{id}.txt")),(os.path.join(to_folder,"pose",f"{test_idx}.txt")))
            # rgb
            shutil.copyfile(
                os.path.join(folder, 'rgb', f'{id}.jpg'),
                os.path.join(to_folder, 'rgb', f'{test_idx}.jpg'),
            )
            # label
            shutil.copyfile(
                os.path.join(folder, 'labels', f'{id}.txt'),
                os.path.join(to_folder, 'labels', f'{test_idx}.txt'),
            )

            # label_images
            shutil.copyfile(
                os.path.join(folder, 'label_images', filename),
                os.path.join(to_folder, 'label_images', f'{test_idx}.jpg'),
            )
            test_idx += 1
